import os
import random
import shutil
import sys
import time
from pathlib import Path

from tqdm import tqdm


def select_labels(src: Path, target: Path, num: int, classes: list[int], strict=True) -> None:
    # Valid parameters
    assert isinstance(src, Path) and isinstance(target, Path), 'ERROR - src or target is not Path'
    assert isinstance(num, int), 'ERROR - num is not int'
    assert isinstance(classes, list), 'ERROR - classes is not list'
    assert all(isinstance(cls, int) for cls in classes), 'ERROR - classes must contain only int(s)'

    # Valid src dir, target dir, num
    assert os.path.isdir(src), 'ERROR - source folder does not exist'
    assert len(os.listdir(src)), 'ERROR - source folder is empty'
    assert num in range(1, len(os.listdir(src))), 'ERROR - invalid num'

    # Create target folder
    create_folder(target)

    # Shuffle labels
    print(f'Shuffling "{src}"...')
    labels = os.listdir(src)
    random.shuffle(labels)

    # Create process bar
    print('Selecting labels...')
    with tqdm(total=num, file=sys.stdout) as pbar:
        for label in labels:
            # The counter
            if not num:
                break

            # Open the source label file and read all valid lines.
            with open(src / label, 'r') as f:
                lines = f.readlines()
                lines = [line.split() for line in lines if int(line[0]) in classes]

            # The new label file is empty and we need exactly `$num` of labels, skip
            if not len(lines) and strict:
                continue

            # Write labels to the new file & save it to target directory
            with open(target / label, 'w') as nf:
                for i in range(len(lines)):
                    nf.write(' '.join(lines[i]))
                    if i != len(lines) - 1:
                        nf.write('\n')

            # The counter
            if len(lines):
                num -= 1
            pbar.update(1)


def select_images(src: Path, target: Path, labels: Path) -> None:
    # Valid parameters
    assert all(isinstance(i, Path) for i in (src, target, labels)), 'ERROR - inputs must be all Path'
    assert os.path.isdir(src), 'ERROR - src does not exist.'
    assert len(os.listdir(src)), 'ERROR - src folder is empty'
    assert os.path.isdir(labels), 'ERROR - label dir does not exist'
    labels = os.listdir(labels)
    assert len(labels), 'ERROR - label dir is empty'

    # Get first image's extension in src folder, and make sure it has valid extension
    extension = os.path.splitext(os.listdir(src)[0])[1]
    assert extension in ('.jpg', '.png'), 'ERROR - src folder must contain .jpg or .png'

    # Create target folder
    create_folder(target)

    print('Selecting images...')
    with tqdm(total=len(labels), file=sys.stdout) as pbar:
        for label in labels:
            image = os.path.splitext(label)[0] + extension
            try:
                shutil.copyfile(src / image, target / image)
            except FileNotFoundError:
                print(f'Failed to copy image "{src / image}"')
                raise
            pbar.update(1)


def select(src_labels: Path, src_images: Path, target_labels: Path, target_images: Path, num: int,
           classes: list[int]) -> None:
    create_folder(Path('.') / 'output')
    select_labels(src_labels, target_labels, num, classes)
    select_images(src_images, target_images, target_labels)
    print('Done!')


def intersect(src_labels: Path, src_images: Path, target_labels: Path, target_images: Path) -> None:
    # Valid parameters
    assert all(isinstance(i, Path) for i in (src_images, src_labels, target_images, target_labels)), \
        'ERROR - all inputs must be Path'

    # Get first image's extension in src folder, and make sure it has valid extension
    extension = os.path.splitext(os.listdir(src_images)[0])[1]
    assert extension in ('.jpg', '.png'), 'ERROR - src folder must contain .jpg or .png'

    # Create target image and label folders
    create_folder(Path('.') / 'dataset', target_images, target_labels)
    intersection = set([os.path.splitext(i)[0] for i in os.listdir(src_labels)]) & set([os.path.splitext(i)[0] for i in
                                                                                        os.listdir(src_images)])

    # All unmatched
    if not intersection:
        raise NotImplementedError(f'No intersection found between "{src_images}" and "{src_labels}"')

    with tqdm(total=len(intersection), file=sys.stdout) as pbar:
        for basename in intersection:
            shutil.copyfile(src_images / (basename + extension), target_images / (basename + extension))
            shutil.copyfile(src_labels / (basename + '.txt'), target_labels / (basename + '.txt'))
            pbar.update(1)
    print('Done!')


def create_folder(*args) -> None:
    for i in args:
        # Delete old folder
        if os.path.exists(i):
            print(f'Deleting old target folder "{i}"...')
            shutil.rmtree(i)

        # Create new folder
        os.makedirs(i, exist_ok=False)
        print(f'Folder "{i}" is created!')


def split(src, target, train_pct=0.7, valid_pct=0.15, test_pct=0.15):
    # Valid parameters
    assert all(isinstance(i, Path) for i in (src, target)), 'ERROR - src & target must be Path'
    labels, images = src / 'labels', src / 'images'
    assert os.path.isdir(labels) and len(os.listdir(labels)), f'ERROR - {labels} must exist and not empty'
    assert os.path.isdir(images) and len(os.listdir(images)), f'ERROR - {images} must exist and not empty'
    assert train_pct + valid_pct + test_pct == 1, f'ERROR - total percentage should be 1'

    labels_li = os.listdir(labels)

    # Create target folder and subfolders
    create_folder(target)
    train = (target / 'train' / 'images', target / 'train' / 'labels')
    valid = (target / 'valid' / 'images', target / 'valid' / 'labels')
    test = (target / 'test' / 'images', target / 'test' / 'labels')
    create_folder(*[item for sublist in (train, valid, test) for item in sublist])

    # Num of files for each subfolder
    total = len(labels_li)
    valid_num, test_num = int(valid_pct * total), int(test_pct * total)

    # Shuffle labels
    print(f'Shuffling "{labels}"...')
    random.shuffle(labels_li)

    print('Splitting labels and images...')
    extension = os.path.splitext(os.listdir(images)[0])[1]
    with tqdm(total=total, file=sys.stdout) as pbar:
        for i, label in enumerate(labels_li):
            basename = os.path.splitext(label)[0]
            if i < valid_num:
                # Copy valid labels & images
                shutil.copyfile(labels / label, valid[1] / label)
                shutil.copyfile(images / (basename + extension), valid[0] / (basename + extension))
            elif i < valid_num + test_num:
                # Copy test labels & images
                shutil.copyfile(labels / label, test[1] / label)
                shutil.copyfile(images / (basename + extension), test[0] / (basename + extension))
            else:
                # Copy train labels & images
                shutil.copyfile(labels / label, train[1] / label)
                shutil.copyfile(images / (basename + extension), train[0] / (basename + extension))
            pbar.update(1)

    print('Done!')


def rename(src: Path, basename: str) -> None:
    # Valid parameter
    assert os.path.isdir(src) and len(os.listdir(src)), f'ERROR - {src} must exist and not empty'
    assert isinstance(basename, str), f'ERROR - {basename} should be str'

    src_li = os.listdir(src)
    extension = os.path.splitext(src_li[0])[1]
    with tqdm(total=len(src_li), file=sys.stdout) as pbar:
        t = int(round(time.time() * 1000))
        for i, c in enumerate(src_li):
            # time.time() to avoid naming conflicts
            os.rename(src / c, src / f'{basename}.{t}.{i}{extension}')
            pbar.update(1)
    print('Done!')
